package com.daxue.common.http.core;

import java.io.Serializable;

/**
 * @author daxue0929
 * @date 2022/4/10
 */

public class TokenBucket implements Serializable {

    private static final long serialVersionUID = 3273853224917588742L;
    private long maxPermits;
    private long intervalMillis;
    private int grantFactor;
    private long backUpTokens;
    private long storedTokens;
    private long nextTokenMillis;

    public TokenBucket init(long now) {
        this.backUpTokens = this.maxPermits;
        this.storedTokens = this.canGranted();
        this.backUpTokens -= this.storedTokens;
        this.nextTokenMillis = now + this.storedTokens * this.intervalMillis;
        return this;
    }

    public long period() {
        return this.maxPermits * this.intervalMillis;
    }

    public TokenBucket reSync(long now) {
        if (now > this.nextTokenMillis + this.intervalMillis) {
            long newTokens = (now - this.nextTokenMillis) / this.intervalMillis;
            this.backUpTokens = Math.min(this.maxPermits - this.storedTokens, this.backUpTokens + newTokens);
            this.nextTokenMillis = Math.addExact(this.nextTokenMillis, newTokens * this.intervalMillis);
        }

        return this;
    }

    public boolean canWait(long now, int tokens, long timeOut) {
        if ((long)tokens > this.maxPermits) {
            return false;
        } else if (timeOut < 0L) {
            return true;
        } else {
            return this.waitToken(now, tokens) < timeOut;
        }
    }

    public long waitToken(long now, int tokens) {
        long needToWait = (long)tokens - Math.min((long)tokens, this.storedTokens);
        long waitMills = 0L;
        long canGranted = 0L;
        if (needToWait > 0L) {
            waitMills = this.nextTokenMillis - now;
            canGranted = this.canGranted();
        }

        if (needToWait > canGranted) {
            waitMills += needToWait * this.intervalMillis;
        }

        return waitMills;
    }

    public long reserveToken(long now, int tokens) {
        long needToWait = (long)tokens - Math.min((long)tokens, this.storedTokens);
        long waitMills = 0L;
        long canGranted = 0L;
        if (needToWait > 0L) {
            waitMills = this.nextTokenMillis - now;
            canGranted = this.canGranted();
        }

        if (needToWait > canGranted) {
            canGranted = needToWait;
        }

        if (needToWait > this.backUpTokens) {
            this.backUpTokens = needToWait;
            waitMills += needToWait * this.intervalMillis;
        }

        this.nextTokenMillis = Math.addExact(this.nextTokenMillis, canGranted * this.intervalMillis);
        this.backUpTokens -= canGranted;
        this.storedTokens += canGranted - (long)tokens;
        return Math.max(waitMills, 0L);
    }

    public long canGranted() {
        return this.grantFactor < 1 ? 0L : (this.backUpTokens + (long)this.grantFactor - 1L) / (long)this.grantFactor;
    }

    private TokenBucket() {
    }

    public static TokenBucket newInstance() {
        return new TokenBucket();
    }

    public long getMaxPermits() {
        return this.maxPermits;
    }

    public long getIntervalMillis() {
        return this.intervalMillis;
    }

    public int getGrantFactor() {
        return this.grantFactor;
    }

    public long getBackUpTokens() {
        return this.backUpTokens;
    }

    public long getStoredTokens() {
        return this.storedTokens;
    }

    public long getNextTokenMillis() {
        return this.nextTokenMillis;
    }

    public TokenBucket setMaxPermits(final long maxPermits) {
        this.maxPermits = maxPermits;
        return this;
    }

    public TokenBucket setIntervalMillis(final long intervalMillis) {
        this.intervalMillis = intervalMillis;
        return this;
    }

    public TokenBucket setGrantFactor(final int grantFactor) {
        this.grantFactor = grantFactor;
        return this;
    }

    public TokenBucket setBackUpTokens(final long backUpTokens) {
        this.backUpTokens = backUpTokens;
        return this;
    }

    public TokenBucket setStoredTokens(final long storedTokens) {
        this.storedTokens = storedTokens;
        return this;
    }

    public TokenBucket setNextTokenMillis(final long nextTokenMillis) {
        this.nextTokenMillis = nextTokenMillis;
        return this;
    }

    public boolean equals(final Object o) {
        if (o == this) {
            return true;
        } else if (!(o instanceof TokenBucket)) {
            return false;
        } else {
            TokenBucket other = (TokenBucket)o;
            if (!other.canEqual(this)) {
                return false;
            } else if (this.getMaxPermits() != other.getMaxPermits()) {
                return false;
            } else if (this.getIntervalMillis() != other.getIntervalMillis()) {
                return false;
            } else if (this.getGrantFactor() != other.getGrantFactor()) {
                return false;
            } else if (this.getBackUpTokens() != other.getBackUpTokens()) {
                return false;
            } else if (this.getStoredTokens() != other.getStoredTokens()) {
                return false;
            } else {
                return this.getNextTokenMillis() == other.getNextTokenMillis();
            }
        }
    }

    protected boolean canEqual(final Object other) {
        return other instanceof TokenBucket;
    }

    public int hashCode() {
        boolean PRIME = true;
        int result = 1;
        long $maxPermits = this.getMaxPermits();
        result = result * 59 + (int)($maxPermits >>> 32 ^ $maxPermits);
        long $intervalMillis = this.getIntervalMillis();
        result = result * 59 + (int)($intervalMillis >>> 32 ^ $intervalMillis);
        result = result * 59 + this.getGrantFactor();
        long $backUpTokens = this.getBackUpTokens();
        result = result * 59 + (int)($backUpTokens >>> 32 ^ $backUpTokens);
        long $storedTokens = this.getStoredTokens();
        result = result * 59 + (int)($storedTokens >>> 32 ^ $storedTokens);
        long $nextTokenMillis = this.getNextTokenMillis();
        result = result * 59 + (int)($nextTokenMillis >>> 32 ^ $nextTokenMillis);
        return result;
    }

    public String toString() {
        return "TokenBucket(maxPermits=" + this.getMaxPermits() + ", intervalMillis=" + this.getIntervalMillis() + ", grantFactor=" + this.getGrantFactor() + ", backUpTokens=" + this.getBackUpTokens() + ", storedTokens=" + this.getStoredTokens() + ", nextTokenMillis=" + this.getNextTokenMillis() + ")";
    }

}
